import numpy as np
import random

class Rfi:
    """
    Class: Rfi is that record and generate rfi various types of data
    Args:
        data(array-like): shape=(NTIMES,NFREQS), rfi amplitude
        flag(array-like): shape=(NTIMES,NFREQS), marked as rfi or not
        freqs:frequency list
        times:time list
    """
    def __init__(self,data,flag,freqs,times):
        self.data = data
        self.flag = flag
        self.freqs = freqs
        self.times = times
    
    def rfi_impulse(self,data_mean,freq=None,time=None,chance=0.05,strength=20):
        """
        Generate a containing RFI impulses that
    are localized in time but span the frequency band.
        Args:
            data_mean: average of background amplitude magnitude
            freq:Add rfi to some specific frequencies,if None,random generate.
            time:Add rfi to some specific times,if None,random generate.
            strength: the amplitude of rfi
        Return: rfidata,rfiflag
        """
        if freq is None:
            count = np.random.randint(0,self.freqs.size / 4)
            freq = np.array(random.sample(list(self.freqs),count))
        if time is None:
            index = np.where(np.random.uniform(size=self.times.size) <= chance)[0]
            # count = np.random.randint(0,self.times.size / 4)
            # time = np.array(random.sample(list(self.times),count))
        # index = np.array([self.times.tolist().index(t) for t in time])
        # dlys = np.random.uniform(-300, 300, size=len(index)) 
        # impulses =  np.array([np.exp(2j * np.pi * np.random.uniform(data_mean,strength*data_mean) * dly * self.freqs) for dly in dlys])
        impulses = np.random.uniform(0,strength*data_mean,size=(len(index),self.freqs.size)) * np.exp(2*np.pi*1j)
        if impulses.size > 0:
            self.data[index] += impulses
            self.flag[index] = 1
        return self.data,self.flag
        
    def rfi_scatter(self,data_mean,chance=0.05,strength=20,std=10):
        """
        Generate a containing scatter RFI 
        Args:
            data_mean: average of background amplitude magnitude
            chance: default=0.05 the probability that a time/freq bin will be assigned an RFI impulse
            strength (float): Jy, default=10
                    the average amplitude of the spike generated in each time/freq bin.
        Return: rfidata,rfiflag
        """
        rfis = np.where(np.random.uniform(size=self.data.size) <= chance)[0]
        # self.data.flat[rfis] += np.random.uniform(data_mean,strength*data_mean) * np.exp(2*np.pi*1j*np.random.uniform(size=rfis.size))
        self.data.flat[rfis] += np.random.uniform(0,strength*data_mean,size=rfis.size) * np.exp(2*np.pi*1j)
        self.flag.flat[rfis] = 1
        return self.data,self.flag

    def rfi_constant(self,data_mean,freq=None,chance=0.05,strength=40,std=10,broadband=False,bandwidth=5):
        """
        Generate a containing constant RFI
        Args:
            data_mean: average of background amplitude magnitude
            freq:Add rfi to some specific frequencies,if None,random generate.
            strength (float): Jy, default=10
                    the average amplitude of the spike generated in each time/freq bin.
            std (float): Jy, default = 10
                    the standard deviation of the amplitudes drawn for each time/freq bin.
            broadband:default=False,Specify whether to generate a broadband signal.
            bandwidth:default=5,representing five frequency points around the center frequency,
             which means that the generated broadband signal has a maximum of five frequency points.
        Return:rfidata,rfiflag
        """
        if freq is None:
            # count = np.random.randint(0,self.freqs.size / 10)
            # freq = np.array(random.sample(list(self.freqs),count))
            index = np.where(np.random.uniform(size=self.freqs.size) <= chance)[0]
            # index = np.array([90,91,92,93,94,95,96,97,98,99,100,101,102, 245, 269, 409, 521, 671,672,673,674, 689, 746, 787, 839, 897, 966])
        if broadband is False:
            # index = np.array([self.freqs.tolist().index(f) for f in freq])
            # freq_samp = np.random.uniform(-300, 300, size=len(index)) 
            # cont = strength * np.array([np.exp(2j * np.pi * fs * self.times) for fs in freq_samp])
            cont = np.random.uniform(0,strength*data_mean,size=(len(index),self.times.size)) * np.exp(2*np.pi*1j)
            self.data.T[index] = cont
            self.flag.T[index] = 1
        else:
            for item in index:
                band = np.random.randint(1,bandwidth)
                sub_band = [] # generate a sub_band
                for i in range(band):
                    sub_band.append(item+i)
                center_cont = np.random.uniform(0,strength*data_mean,size=self.times.size) * np.exp(2*np.pi*1j)
                sub_cont = []
                for i in range(band // 2 - 1,-1,-1):
                    sub_cont.append(center_cont / 2**i)
                sub_cont.append(center_cont)
                for i in range(band // 2 + 1,band):
                    sub_cont.append(center_cont / 2**i)
                self.data.T[sub_band] = sub_cont
                self.flag.T[sub_band] = 1

        return self.data,self.flag
    def rfi_telescope_style(self,filename):
        """
        """